#!/usr/bin/env python3
"""
修正版：mpc_quad_gpt1.1_fixed.py
修复：
 - 用到最终目标 TARGET_POS 判断是否到达（而不是当前参考点 x_ref）
 - 补上 pd_hover_control 定义（用于悬停阶段）
 - 改善绘图在单点时的 xlim/ylim 设置，避免 warnings
 - 其余逻辑与你之前运行的大体相同（权重/horizon/waypoint/scaling 优化等）
"""
import os
import time as pytime
from math import ceil
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ---------------------------
# PARAMETERS (和你之前脚本保持一致，可调整)
# ---------------------------
DT = 0.1
T_TOTAL = 60.0
STEPS = int(ceil(T_TOTAL / DT))

# tuning flags (和你之前使用的一致)
APPLY_WEIGHT_TUNING = True
R_FACTOR = 10.0
Q_VEL_FACTOR = 5.0
APPLY_SHORT_HORIZON = True
HORIZON = 50
SHORT_HORIZON = 20
APPLY_WAYPOINT = True
WAYPOINT_MAX_SPEED = 5.0
USE_SCALING = True

# Actuator limits
U_MIN = np.array([-5.0, -5.0, -5.0])
U_MAX = np.array([ 5.0,  5.0,  5.0])

# Target and initial
TARGET_POS = np.array([80.0, -20.0, 30.0])
X0 = np.zeros(6); X0[2] = 0.0

# tolerances
POS_TOL = 0.5
VEL_TOL = 0.5
HOVER_TIME = 8.0
PRED_SHOW = 12

# ---------------------------
# Dynamics & cost
# ---------------------------
I3 = np.eye(3)
A = np.block([[I3, DT * I3],
              [np.zeros((3,3)), I3]])
B = np.block([[0.5 * DT**2 * I3],
              [DT * I3]])

Q_pos = 10.0
Q_vel = 1.0
R_base = 0.1 * np.eye(3)

# tuned
Q_vel_tuned = Q_vel * Q_VEL_FACTOR if APPLY_WEIGHT_TUNING else Q_vel
R_tuned = R_base * R_FACTOR if APPLY_WEIGHT_TUNING else R_base
horizon = SHORT_HORIZON if APPLY_SHORT_HORIZON else HORIZON
Q = np.diag([Q_pos, Q_pos, Q_pos, Q_vel_tuned, Q_vel_tuned, Q_vel_tuned])
R = R_tuned

# ---------------------------
# finite-horizon LQR gains
# ---------------------------
def finite_horizon_lqr_gains(A,B,Q,R,N):
    P = [None]*(N+1)
    P[N] = Q.copy()
    K = [None]*N
    for k in range(N-1, -1, -1):
        Bt_PB = B.T @ P[k+1] @ B
        inv = np.linalg.inv(R + Bt_PB)
        Kk = inv @ (B.T @ P[k+1] @ A)
        K[k] = Kk
        Pk = Q + A.T @ P[k+1] @ A - A.T @ P[k+1] @ B @ Kk
        P[k] = Pk
    return K, P

K_seq, P_seq = finite_horizon_lqr_gains(A,B,Q,R,horizon)
K0 = K_seq[0]

# ---------------------------
# Reference generation (speed-limited)
# ---------------------------
def generate_speed_limited_ref(x_start, target_pos, dt, max_speed, steps):
    start_pos = x_start[0:3].copy()
    vec = target_pos - start_pos
    dist = np.linalg.norm(vec)
    if dist < 1e-6:
        return np.tile(target_pos, (steps,1))
    dir_unit = vec / dist
    max_step = max_speed * dt
    n_steps_needed = int(np.ceil((dist - 0.5) / max_step)) if dist > 0.5 else 1
    n_steps_needed = max(1, n_steps_needed)
    pos_seq = [start_pos + dir_unit * min(max_step * i, dist) for i in range(n_steps_needed+1)]
    if len(pos_seq) < steps:
        pad = [target_pos] * (steps - len(pos_seq))
        pos_seq.extend(pad)
    pos_seq = np.array(pos_seq[:steps])
    return pos_seq

if APPLY_WAYPOINT:
    ref_pos_seq = generate_speed_limited_ref(X0, TARGET_POS, DT, WAYPOINT_MAX_SPEED, STEPS)
else:
    ref_pos_seq = np.tile(TARGET_POS, (STEPS,1))

def ref_state_at(step_idx):
    p = ref_pos_seq[min(step_idx, ref_pos_seq.shape[0]-1)]
    xr = np.zeros(6); xr[0:3] = p; xr[3:6] = 0.0
    return xr

# ---------------------------
# Prediction helper
# ---------------------------
def predict_trajectory(x0, x_ref, K_seq, A, B, steps):
    x = x0.copy()
    preds = [x.copy()]
    u_preds = []
    for k in range(min(steps, len(K_seq))):
        err = x - x_ref
        u = -K_seq[k].dot(err)
        u_preds.append(u.copy())
        x = A.dot(x) + B.dot(u)
        preds.append(x.copy())
    return np.array(preds), np.array(u_preds)

# ---------------------------
# PD hover control (补上缺失函数)
# ---------------------------
def pd_hover_control(pos, vel, target_pos, kp=0.6, kd=1.0):
    """简单 PD，用于悬停保持"""
    pos_err = pos - target_pos
    u = -kp * pos_err - kd * vel
    return u

# ---------------------------
# scaling util
# ---------------------------
def scale_to_limits(u_vec, u_min, u_max):
    max_comp = np.max(np.abs(u_vec))
    max_limit = np.max(np.abs(u_max))
    if max_comp <= max_limit or max_comp == 0.0:
        return u_vec.copy(), False, 1.0
    factor = max_limit / max_comp
    return u_vec * factor, True, factor

# ---------------------------
# plotting setup
# ---------------------------
out_dir = os.path.dirname(os.path.abspath(__file__))
log_path = os.path.join(out_dir, 'mpc_quad_opt_log_fixed.csv')
png_path = os.path.join(out_dir, 'mpc_quad_opt_fixed.png')

times = []; x_hist = []; u_hist = []; u_raw_hist = []; scaled_hist = []

x = X0.copy()
x_ref = ref_state_at(0)

interactive_ok = True
fig = None
try:
    plt.ion()
    fig = plt.figure(figsize=(14,8))
    gs = fig.add_gridspec(2, 3, width_ratios=[2.8,1,1])
    ax3d = fig.add_subplot(gs[:,0], projection='3d')
    ax_pos = fig.add_subplot(gs[0,1])
    ax_ctrl = fig.add_subplot(gs[1,1])
    ax_u_raw = fig.add_subplot(gs[0,2])
    ax_info = fig.add_subplot(gs[1,2]); ax_info.axis('off')
    info_text = ax_info.text(0.01, 0.98, "", va='top', ha='left', fontsize=9, family='monospace')
    traj_line, = ax3d.plot([], [], [], lw=2, label='trajectory')
    start_scatter = ax3d.scatter([X0[0]],[X0[1]],[X0[2]], c='b', s=80, marker='o', label='start')
    target_scatter = ax3d.scatter([TARGET_POS[0]],[TARGET_POS[1]],[TARGET_POS[2]], c='r', s=120, marker='x', label='target')
    pred_line, = ax3d.plot([], [], [], ls='--', lw=1.3, color='orange', label='pred_traj')
    pred_quiver = None
    final_scatter = ax3d.scatter([], [], [], c='g', s=150, marker='*', label='final')
    ax3d.set_xlabel('x (m)'); ax3d.set_ylabel('y (m)'); ax3d.set_zlabel('z (m)')
    ax3d.set_title('3D Trajectory (start blue, target red, pred orange)')
    ax3d.legend(loc='upper left')

    pos_lines = [ax_pos.plot([], [], label='x')[0], ax_pos.plot([], [], label='y')[0], ax_pos.plot([], [], label='z')[0]]
    ax_pos.set_xlabel('time (s)'); ax_pos.set_ylabel('position (m)'); ax_pos.set_title('Position vs time'); ax_pos.legend(); ax_pos.grid(True)

    ctrl_lines = [ax_ctrl.plot([], [], label='ax')[0], ax_ctrl.plot([], [], label='ay')[0], ax_ctrl.plot([], [], label='az')[0]]
    ax_ctrl.set_xlabel('time (s)'); ax_ctrl.set_ylabel('acc (m/s^2)'); ax_ctrl.set_title('Applied control u vs time'); ax_ctrl.legend(); ax_ctrl.grid(True)

    uraw_lines = [ax_u_raw.plot([], [], label='u_raw_x')[0],
                  ax_u_raw.plot([], [], label='u_raw_y')[0],
                  ax_u_raw.plot([], [], label='u_raw_z')[0]]
    uapp_lines = [ax_u_raw.plot([], [], ls='--', label='u_x')[0],
                  ax_u_raw.plot([], [], ls='--', label='u_y')[0],
                  ax_u_raw.plot([], [], ls='--', label='u_z')[0]]
    ax_u_raw.set_title('u_raw solid / applied u dashed ; red dots = scaling/clipping events')
    ax_u_raw.set_xlabel('time (s)'); ax_u_raw.set_ylabel('acc (m/s^2)'); ax_u_raw.legend(); ax_u_raw.grid(True)

    plt.tight_layout()
except Exception as e:
    print("Interactive plotting unavailable; will save static final images. Exception:", e)
    interactive_ok = False
    try:
        plt.ioff()
    except:
        pass

def remove_quiver(handle):
    try:
        if handle is not None:
            handle.remove()
    except Exception:
        pass

def update_plots(t_arr, X_arr, U_arr, Uraw_arr, preds_pos=None, preds_vel=None, scaled_flags=None, info_str=""):
    global interactive_ok, pred_quiver
    if not interactive_ok or X_arr.shape[0] == 0:
        return
    try:
        # set data
        traj_line.set_data(X_arr[:,0], X_arr[:,1]); traj_line.set_3d_properties(X_arr[:,2])
        if preds_pos is not None and preds_pos.shape[0] > 0:
            pred_line.set_data(preds_pos[:,0], preds_pos[:,1]); pred_line.set_3d_properties(preds_pos[:,2])
        else:
            pred_line.set_data([], []); pred_line.set_3d_properties([])

        remove_quiver(pred_quiver); pred_quiver = None
        if preds_pos is not None and preds_pos.shape[0] > 1:
            vel_scale = 0.2
            Xp = preds_pos[:,0]; Yp = preds_pos[:,1]; Zp = preds_pos[:,2]
            Up = preds_vel[:,0]*vel_scale; Vp = preds_vel[:,1]*vel_scale; Wp = preds_vel[:,2]*vel_scale
            try:
                pred_quiver = ax3d.quiver(Xp, Yp, Zp, Up, Vp, Wp, length=1.0, normalize=False, color='orange', linewidth=1)
            except Exception:
                for i in range(len(Xp)):
                    ax3d.plot([Xp[i], Xp[i]+Up[i]], [Yp[i], Yp[i]+Vp[i]], [Zp[i], Zp[i]+Wp[i]], color='orange', alpha=0.8)

        # safe limits when only 1 time point
        if t_arr.size == 1:
            t0 = float(t_arr[0])
            ax_pos.set_xlim(t0-0.1, t0+0.1)
            ax_ctrl.set_xlim(t0-0.1, t0+0.1)
            ax_u_raw.set_xlim(t0-0.1, t0+0.1)
        else:
            ax_pos.set_xlim(t_arr[0], t_arr[-1])
            ax_ctrl.set_xlim(t_arr[0], t_arr[-1])
            ax_u_raw.set_xlim(t_arr[0], t_arr[-1])

        # pos/ctrl
        for i, ln in enumerate(pos_lines):
            ln.set_data(t_arr, X_arr[:, i])
        all_pos = X_arr[:,0:3].flatten(); ax_pos.set_ylim(np.min(all_pos)-5, np.max(all_pos)+5)

        for i, ln in enumerate(ctrl_lines):
            ln.set_data(t_arr, U_arr[:, i])
        all_ctrl = U_arr.flatten(); ax_ctrl.set_ylim(np.min(all_ctrl)-1.0, np.max(all_ctrl)+1.0)

        for i, ln in enumerate(uraw_lines):
            ln.set_data(t_arr, Uraw_arr[:, i])
        for i, ln in enumerate(uapp_lines):
            ln.set_data(t_arr, U_arr[:, i])
        all_uraw = Uraw_arr.flatten(); ax_u_raw.set_ylim(np.min(all_uraw)-5.0, np.max(all_uraw)+5.0)

        # mark scaled/clipped events
        if scaled_flags is not None and scaled_flags.size == t_arr.size:
            times_mask = np.array(scaled_flags, dtype=bool)
            if np.any(times_mask):
                tsc = t_arr[times_mask]
                for axis in range(3):
                    vals = Uraw_arr[:, axis]
                    ax_u_raw.scatter(tsc, vals[times_mask], c='r', s=10)

        info_text.set_text(info_str)
        fig.canvas.draw(); fig.canvas.flush_events()
    except Exception as e:
        print("Plot update error:", e)
        raise

# ---------------------------
# Main loop (修复到达判定 & 使用 pd_hover_control)
# ---------------------------
print("Running fixed version. settings approx same as before.")
sim_start = pytime.time()

hover_mode = False; hover_start = None; reached = False

for step in range(STEPS):
    t = step * DT
    # current reference (time-varying if waypoint)
    x_ref = ref_state_at(step)
    err_ref = x - x_ref
    u_raw = -K0.dot(err_ref)

    preds, u_preds = predict_trajectory(x, x_ref, K_seq, A, B, PRED_SHOW)
    preds_pos = preds[:,0:3] if preds is not None else None
    preds_vel = preds[:,3:6] if preds is not None else None

    # scaling
    scaled_flag = False; scale_factor = 1.0
    u_to_apply = u_raw.copy()
    if not hover_mode:
        if USE_SCALING:
            u_scaled, scaled_flag, scale_factor = scale_to_limits(u_raw, U_MIN, U_MAX)
            u_to_apply = u_scaled
        u = np.maximum(U_MIN, np.minimum(U_MAX, u_to_apply))
    else:
        pos = x[0:3]; vel = x[3:6]
        u_des = pd_hover_control(pos, vel, TARGET_POS, kp=0.6, kd=1.0)
        u = np.maximum(U_MIN, np.minimum(U_MAX, u_des))
        scaled_flag = False; scale_factor = 1.0

    # dynamics (small/no disturbance)
    w = (np.random.randn(3)*0.0)
    x = A.dot(x) + B.dot(u) + B.dot(w)

    # logs
    times.append(t); x_hist.append(x.copy()); u_hist.append(u.copy()); u_raw_hist.append(u_raw.copy()); scaled_hist.append(scaled_flag)

    # diagnostics
    pos = x[0:3]; vel = x[3:6]
    err_to_target = np.linalg.norm(pos - TARGET_POS)   # <-- 用到最终目标判断到达
    err_ref_norm = np.linalg.norm(err_ref[0:3])
    sat_axes = int(np.sum((u_raw < U_MIN) | (u_raw > U_MAX)))
    u_raw_norm = np.linalg.norm(u_raw); u_norm = np.linalg.norm(u)
    pred_final = preds_pos[-1] if preds_pos is not None else None
    pred_final_dist = np.linalg.norm(pred_final - TARGET_POS) if pred_final is not None else None

    # info string
    info_lines = [
        f"t={t:5.2f}s  mode={'HOVER' if hover_mode else 'MPC'}  scaled={scaled_flag} sf={scale_factor:.3f}",
        f"pos=({pos[0]:6.1f},{pos[1]:6.1f},{pos[2]:5.1f})  vel=({vel[0]:5.1f},{vel[1]:5.1f},{vel[2]:5.1f})",
        f"u_raw=({u_raw[0]:7.2f},{u_raw[1]:7.2f},{u_raw[2]:7.2f}) -> applied u=({u[0]:6.2f},{u[1]:6.2f},{u[2]:6.2f})",
        f"||u_raw||={u_raw_norm:7.2f} ||u||={u_norm:6.2f}  sat_axes={sat_axes}/3  err_to_target={err_to_target:6.2f}m",
    ]
    if pred_final is not None:
        info_lines.append(f"pred_final (h={PRED_SHOW}) = ({pred_final[0]:6.1f},{pred_final[1]:6.1f},{pred_final[2]:5.1f}) dist_pred={pred_final_dist:6.2f}")
    info_lines.append("")
    info_lines.append("Key: arrival check uses distance to FINAL TARGET (not to moving reference).")
    info_str = "\n".join(info_lines)

    # console prints
    print(f"[t={t:5.2f}] pos=({pos[0]:6.1f},{pos[1]:6.1f},{pos[2]:5.1f}) vel=({vel[0]:5.2f},{vel[1]:5.2f},{vel[2]:5.2f})")
    print(f"    u_raw=({u_raw[0]:7.2f},{u_raw[1]:7.2f},{u_raw[2]:7.2f}) -> applied u=({u[0]:6.2f},{u[1]:6.2f},{u[2]:6.2f}) scaled={scaled_flag} sat_axes={sat_axes}")
    if preds_pos is not None:
        showN = min(PRED_SHOW, preds_pos.shape[0]-1)
        pred_list = " ".join([f"({preds_pos[i+1,0]:5.1f},{preds_pos[i+1,1]:5.1f},{preds_pos[i+1,2]:4.1f})" for i in range(showN)])
        print(f"    Pred next {showN}: {pred_list}  pred_final_dist={pred_final_dist:6.2f}")

    # ARRIVAL: 使用到最终目标的误差和速度判断（避免 t=0 的误判）
    if (err_to_target < POS_TOL) and (np.linalg.norm(vel) < VEL_TOL) and (not hover_mode):
        hover_mode = True; hover_start = t
        print(f"*** ARRIVED at FINAL TARGET -> entering HOVER for {HOVER_TIME} s (t={t:.2f}) ***")

    if hover_mode and (t - hover_start >= HOVER_TIME):
        print("*** HOVER complete -> ending sim ***"); reached = True; break

    # update plots
    try:
        t_arr = np.array(times)
        X_arr = np.array(x_hist)
        U_arr = np.array(u_hist)
        Uraw_arr = np.array(u_raw_hist)
        scaled_flags = np.array(scaled_hist)
        update_plots(t_arr, X_arr, U_arr, Uraw_arr, preds_pos, preds_vel, scaled_flags, info_str)
        if interactive_ok:
            plt.pause(0.001)
    except KeyboardInterrupt:
        print("Interrupted by user"); break
    except Exception as e:
        print("Plot update error:", e); break

sim_end = pytime.time()
print(f"Simulation ended. wall_time={sim_end-sim_start:.2f}s steps_simulated={len(times)} reached={reached}")

# Save CSV
Xnp = np.array(x_hist); Unp = np.array(u_hist); Urawnp = np.array(u_raw_hist)
df = pd.DataFrame({
    't': np.array(times),
    'x': Xnp[:,0], 'y': Xnp[:,1], 'z': Xnp[:,2],
    'vx': Xnp[:,3], 'vy': Xnp[:,4], 'vz': Xnp[:,5],
    'ax': Unp[:,0], 'ay': Unp[:,1], 'az': Unp[:,2],
    'u_raw_x': Urawnp[:,0], 'u_raw_y': Urawnp[:,1], 'u_raw_z': Urawnp[:,2],
    'scaled_flag': np.array(scaled_hist, dtype=int)
})
df.to_csv(log_path, index=False)
print("Saved CSV log to:", log_path)

# Save final figure
try:
    final_pos = Xnp[-1, 0:3]
    if interactive_ok and fig is not None:
        ax3d.scatter([final_pos[0]],[final_pos[1]],[final_pos[2]], c='g', s=150, marker='*', label='final')
        ax3d.legend(); fig.savefig(png_path, dpi=160, bbox_inches='tight')
    else:
        fig_s = plt.figure(figsize=(14,8)); gs2 = fig_s.add_gridspec(2,3, width_ratios=[2.8,1,1])
        a3 = fig_s.add_subplot(gs2[:,0], projection='3d'); ap = fig_s.add_subplot(gs2[0,1]); ac = fig_s.add_subplot(gs2[1,1]); au = fig_s.add_subplot(gs2[0,2])
        a3.plot(Xnp[:,0], Xnp[:,1], Xnp[:,2], lw=2); a3.scatter([X0[0]],[X0[1]],[X0[2]], c='b', s=80, marker='o', label='start')
        a3.scatter([TARGET_POS[0]],[TARGET_POS[1]],[TARGET_POS[2]], c='r', s=120, marker='x', label='target')
        a3.scatter([final_pos[0]],[final_pos[1]],[final_pos[2]], c='g', s=150, marker='*', label='final')
        a3.set_title('3D Trajectory (final)'); a3.set_xlabel('x'); a3.set_ylabel('y'); a3.set_zlabel('z')
        ap.plot(df['t'], df['x'], label='x'); ap.plot(df['t'], df['y'], label='y'); ap.plot(df['t'], df['z'], label='z'); ap.legend(); ap.grid(True)
        ac.plot(df['t'], df['ax'], label='ax'); ac.plot(df['t'], df['ay'], label='ay'); ac.plot(df['t'], df['az'], label='az'); ac.legend(); ac.grid(True)
        au.plot(df['t'], df['u_raw_x'], label='u_raw_x'); au.plot(df['t'], df['u_raw_y'], label='u_raw_y'); au.plot(df['t'], df['u_raw_z'], label='u_raw_z'); au.legend(); au.grid(True)
        fig_s.tight_layout(); fig_s.savefig(png_path, dpi=160, bbox_inches='tight')
    print("Saved combined figure to:", png_path)
except Exception as e:
    print("Failed to save combined figure:", e)

if interactive_ok:
    try:
        print("Interactive: close figure window to finish."); plt.ioff(); plt.show()
    except:
        pass
